import cv2
import os

from mmdet.apis import init_detector, inference_detector, show_result_pyplot
from pycocotools.coco import COCO
config_file = '/home/ubuntu/code/mmdetection-master/pollen_experiment/configs/mask_rcnn_config.py'
# 从 model zoo 下载 checkpoint 并放在 `checkpoints/` 文件下
# 网址为: http://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth
checkpoint_file = '/home/ubuntu/code/mmdetection-master/result/mask-rcnn/latest.pth'
device = 'cuda:0'
# 初始化检测器
model = init_detector(config_file, checkpoint_file, device=device)
model.eval()
# 推理演示图像
# img = '/home/ubuntu/code/detectron2-latest/datasets/pollen_data/0318/test/JPEGImages/186_89.jpg'

if hasattr(model, 'module'):
    model = model.module
coco = COCO('/home/ubuntu/code/mmdetection-master/pollen_data/paper_test.json')
images_path = '/home/ubuntu/code/detectron2-latest/datasets/pollen_data/0318/imgs'
imgs = coco.loadImgs(coco.getImgIds())

for i in imgs:
    img = os.path.join(images_path, i['file_name'])
    result = inference_detector(model, img)
    predict_img = model.show_result(
        img,
        result,
        score_thr=0.5,
        show=False,
        wait_time=0,
        win_name='result',
        bbox_color=None,
        text_color=(200, 200, 200),
        mask_color=None)
    cv2.imwrite(os.path.join('/home/ubuntu/code/mmdetection-master/result/mask-rcnn/predict_images', i['file_name']), predict_img)
